package com.suieqian.config.dblimit;

import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.extension.parser.JsqlParserSupport;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import lombok.Setter;
import net.sf.jsqlparser.expression.Function;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectExpressionItem;
import net.sf.jsqlparser.statement.select.SelectItem;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.sql.SQLException;
import java.util.Objects;

@Setter
public class DbRecordLimitInterceptor extends JsqlParserSupport implements InnerInterceptor {
    private int limitSize = 1000;

    /**
     *
     * <p>
     * 改改sql啥的
     *
     * @param executor      Executor(可能是代理对象)
     * @param ms            MappedStatement
     * @param parameter     parameter
     * @param rowBounds     rowBounds
     * @param resultHandler resultHandler
     * @param boundSql      boundSql
     */
    @Override
    public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
        PluginUtils.MPBoundSql mpBs = PluginUtils.mpBoundSql(boundSql);
        try {
            Statement stmt = CCJSqlParserUtil.parse(boundSql.getSql());
            if(!(stmt instanceof Select)){
                return;
            }
            PlainSelect select = (PlainSelect)((Select) stmt).getSelectBody();
            //已有limit
            if(Objects.nonNull(select.getLimit())){
                return;
            }
            for (SelectItem item :select.getSelectItems()) {
                if (item.toString().equals("*")) {
                    return;
                }
                SelectExpressionItem exItem = (SelectExpressionItem) item;
                //带着函数，eg：count，max，avg，sum
                if(exItem.getExpression() instanceof Function){
                    return;
                }
            }
            //加入limit
            mpBs.sql(boundSql.getSql()+" limit "+limitSize);
        } catch (Exception e) {
            e.printStackTrace();
        }

    }
}
